import warnings
warnings.filterwarnings("ignore")
import os
os.environ["TF_CPP_MIN_LOG_LEVEL"]="2"
import tensorflow as tf
tf.compat.v1.logging.set_verbosity(40)

import numpy as np
from scipy.io import wavfile
import os
import scipy.io.wavfile
from tensorflow.keras import losses,optimizers,applications
from tensorflow.keras.layers import Dense,Input
from tensorflow.keras.models import Model
from scipy import signal
from sklearn.model_selection import train_test_split

#得到取对数之后的语谱图
def get_log_spec(audio,sr,window_length=20,step_length=10,eps=1e-8):
    nperseg=int(round(window_length*sr/1e3))#帧长  160
    noverlap=int(round(step_length*sr/1e3))#重合部分  80
    freq,time,spec=signal.spectrogram(x=audio,
                                      fs=sr,
                                      nperseg=nperseg,
                                      noverlap=noverlap,
                                      detrend=False)
    return freq,time,np.log(spec.T.astype(np.float32)+eps)
'''
    x---输入信号的向量
    fs---采样频率Hz，如果指定为[]，默认为1Hz
    window---窗函数，如果window为一个整数，x将被分成window段，每段使用Hamming窗函数加窗。
             如果window是一个向量，x将被分成length(window)段，每一段使用window向量指定的窗函数加窗。
             所以如果想获取specgram函数的功能，只需指定一个256长度的Hann窗
    nperseg:int型 窗口长度
    noverlap---各段之间重叠的采样点数
    detrend : str 或 function 或 False，可选指定如何去趋势每个段。 
    如果 detrend 是一个字符串，它将作为类型参数传递给 detrend 函数。 
    如果它是一个函数，它需要一个段并返回一个去趋势段。 
    如果去趋势为 False，则不进行去趋势。 默认为“常量”
    返回值：f:频域信息  t:时域信息  Sxx:时频图(语谱图)  这里相当于取log之后的filter_bank  目前还没有得到mfcc
'''

#读取数据
def readData(path):
    x_data=[]
    y_data=[]
    path_data=[]
    new_sr=8000
    for i,class_name in enumerate(os.listdir(path)):
        class_path=os.path.join(path,class_name)
        for voice_name in os.listdir(class_path):
            voice_path=os.path.join(class_path,voice_name)
            sr,sample=scipy.io.wavfile.read(voice_path)#sr 采样率   sample 音频信号值
            sample=sample[:,0]#取单声道
            resample=signal.resample(sample,int(new_sr/sr*sample.shape[0]))#按照sr的改变比率  来控制sample的改变  simple_shape:(16000,)  resample_shape:(8000,)
            _,_,spec=get_log_spec(resample,new_sr)#得到取对数之后的语谱图
            x_data.append(spec)
            y_data.append(np.eye(3)[i])#独热
            path_data.append(voice_path)
    return np.array(x_data),np.array(y_data),np.array(path_data)

#采用轻量级的mobilnet模型
def my_cnn(input_shape,num_class=3):
    input=Input(input_shape)
    pre_train_model=applications.MobileNet(input_shape=input_shape,
              include_top=False,
              pooling='avg')
    pre_train_model.trainable=True
    x=pre_train_model(input)
    output=Dense(num_class,activation='softmax')(x)
    model=Model(input,output)
    return model

if __name__ == '__main__':
    #读取数据
    data_path=r"D:\data_all\project-voice\voice_data_3\train"
    lable_names = os.listdir(data_path)
    x_data,y_data,path_data=readData(data_path)
    x_data=np.stack((x_data,x_data,x_data),axis=3)
    print(x_data.shape)

    #划分数据集
    x_trian,x_test,y_trian,y_test,path_train,path_test=train_test_split(x_data,y_data,path_data,test_size=0.2)
    x_test,x_val,y_test,y_val,path_test,path_val=train_test_split(x_test,y_test,path_test,test_size=0.5)

    #创建模型，进行训练
    model=my_cnn(x_data.shape[1:],len(lable_names))
    model.summary()
    model.compile(loss=losses.categorical_crossentropy,
                  optimizer=optimizers.Adam(),
                  metrics=['acc'])
    history = model.fit(x_trian,y_trian, epochs=10, batch_size=8, validation_data=(x_val, y_val))

    #保存模型为.h5文件
    model.save('MobileNet_3.h5')

    # converter = tf.lite.TFLiteConverter.from_keras_model(model)
    # tflite_model = converter.convert()
    # open("MobileNet.tflite", "wb").write(tflite_model)
    #
    # score = model.evaluate(x_test, y_test, batch_size=8)
    # # model.save("MobileNet.h5")
    # print("loss_test:", score[0])
    # print("loss_accuracy:", score[1])


